#!/usr/bin/env python
# encoding: UTF-8

"""Finds interfaces in all files for which the contacts tables in
the db have been generated. A new table, interfaces, is created.
If there was such a table in the db, it is removed and recreated prior
to adding any records."""

from Bio.PDB.PDBParser import PDBParser
from db_cleanup import get_tables
from itertools import tee, takewhile

import sys
import sqlite3
import os.path
import logging
import warnings

log = logging.getLogger("find_interfaces")
log.addHandler(logging.StreamHandler())
log.setLevel(logging.INFO)

def main(argv):
    try:
        db_name = argv[1]
        pdb_dir = argv[2]
    except IndexError:
        print "Usage: %s database pdb_directory" % argv[0]
        print
        print __doc__
        return 1
        
    pdbs = [pdb.split("_")[1] for (pdb, _) in get_tables(db_name)]
    count = len(pdbs)
    
    p = PDBParser()
    
    conn = sqlite3.connect(db_name)
    cur = conn.cursor()
    cur2 = conn.cursor()
    
    cur.executescript("""
    DROP TABLE IF EXISTS interfaces;
    CREATE TABLE interfaces (
        "id" integer PRIMARY KEY,
        "pdb" char(4),
        "chain" char(2),
        "start_residue" integer,
        "end_residue" integer,
        "num_residues" integer);""")

    problems = []

    for i, pdb in enumerate(pdbs):
        log.info("%s (%d/%d)" % (pdb, i+1, count))
        try:
            model = p.get_structure(pdb, os.path.join(pdb_dir, pdb + ".pdb"))[0]
        except IOError, e:
            log.warn("Cannot process structure %s: %s" % (pdb, e))
            problems.append(pdb)
            continue
        for chain in model:
            i = find_interfacing_residues(chain, pdb, cur)
            i1, i2 = tee(i)
            resids = [x[0] for x in i1]
            interfacing = list(x[1] for x in i2)

            for interface in make_interfaces(interfacing, resids):
                start = interface[0]
                end = interface[1]
                count_res = interface[2]
                cur2.execute("""INSERT INTO interfaces (pdb, chain, start_residue, end_residue, num_residues)
                                                VALUES(?,?,?,?,?)""", (pdb, chain.id, start, end, count_res))
        conn.commit()
    cur2.execute("VACUUM")
    conn.close()
    if len(problems) == 0:
        log.info("All contacts in the database were processed.")
    else:
        log.info("Some data could not be processed: %s" % problems)
    return 0

def make_interfaces(interfacing, residues):
    """Takes a list of boolean values and residue ids. The truth value
    in the first list indicates whether the residue id from the second
    is in contact with other chain.

    Returns iterator of tuples indicating the interfaces in the form
    (start_id, end_id, count_of_contacts).

    The first residue of an interface is the first residue in contact,
    the same for the end of the interface. The interface continues
    through one residue gaps. The lone one-residue contacts are not
    considered an interface.

    Example data (1 = True, 0 = False for brevity):
    interfacing: 01101100010101100111
    residues:    01234567890123456789 (the second half is +10)
    Resulting interfaces:
        1,2  2
        4,5  2
       11,14 3
       17,19 3
    """
    if len(interfacing) != len(residues):
        raise ValueError("Length of 'interfacing' and 'residues' must be the same")
    resids = residues
    start = None
    count = 0
    interfaces = []
    
    for i, is_interfacing in enumerate(interfacing):
        if is_interfacing:
            count += 1
            if start is None:
                start = resids[i]
        else:
            if start is None:
                continue
            try:
                if interfacing[i+1]:
                    continue # as we will still be in the interface in the next step
                else: # end of an interface
                    interfaces.append((start, resids[i-1], count))
                    count = 0
                    start = None
            except IndexError:
                continue # end of the chain, the interface will be added in the next step

    if start is not None:
        # add the interface if it continues to the last residue
        interfaces.append((start, resids[-1], count)) 

    interfaces = (x for x in interfaces if x[1] != x[0])
    return interfaces
            
            
            
            
def find_interfacing_residues(chain, pdb, cur):
    """The contact is at least one hydrogen bond (direct or wm)
    or direct van der Waals contact.

    Returns iterator of tuples in the form (residue_id, is_in_contact).
    """
    N_CONTACTS = 1 # at least one contact per interface residue is required
    for residue in chain:
        het, number, insertion = residue.get_id()
        name = residue.get_resname().strip()
        if name == "HOH":
            continue

        cur.execute("""SELECT a.chain, a.residue, a.amino, d.chain, d.residue, d.amino
        FROM contacts_%(pdb)s c, atoms_%(pdb)s a, atoms_%(pdb)s d
        WHERE donor_id = d.id
            AND acceptor_id = a.id
            AND (type = "H" OR directness = "D")
            AND (  (a.chain = ?
                    AND a.residue = ?)
                OR (
                    d.chain = ?
                    AND d.residue = ?)
                )""" % {"pdb":pdb}, (chain.id, number)*2) # maybe the chain constraint is redundant?
                
        cs = cur.fetchall()
        
        in_interface = True if len(cs) >= N_CONTACTS else False
        yield number, in_interface
                
            

if __name__ == "__main__":
    sys.exit(main(sys.argv))
